import torch
import sys

sys.path.append(r'D:\mydata\heima\heimaPY\KG\LSTM_CRF\utils')
from common import *
from torch.utils.data import DataLoader, Dataset
from torch.nn.utils.rnn import pad_sequence
datas, word2id = build_data()
# 构建Dataset类
class NerDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, item):
        x = self.data[item][0]
        y = self.data[item][1]
        return x, y


def collate_fn(batch):
    # 获取数据x、标签y
    x = [torch.tensor([word2id[word] for word in sample[0]]) for sample in batch]
    y = [torch.tensor([config.tag2id[tag] for tag in sample[1]]) for sample in batch]
    # 长度补齐
    x = pad_sequence(x, batch_first=True, padding_value=0)
    # BiLSTM用-100补齐，BiLSTM_CRF用0补齐
    if config.model == "BiLSTM":
        y = pad_sequence(y, batch_first=True, padding_value=-100)
    elif config.model == "BiLSTM_CRF":
        y = pad_sequence(y, batch_first=True, padding_value=-0)
    # 获得mask
    mask = (x != 0).long()
    return x, y, mask


# 获取dataloader
def get_data():
    """
    获取 dataloder 数据迭代器
    :return:train_dataloader,test_dataloader
            训练迭代器，测试迭代器
    """
    # 1、获取训练集dataloader
    # dataset
    train_dataset = NerDataset(datas[:6200])
    # dataloader
    train_dataloader = DataLoader(dataset=train_dataset,
                                  batch_size=config.batch_size,
                                  collate_fn=collate_fn,
                                  drop_last=True,
                                  shuffle=True)

    # 2、获取测试集dataloader
    # dataset
    test_dataset = NerDataset(datas[6200:])
    # dataloader
    test_dataloader = DataLoader(dataset=test_dataset,
                                 batch_size=config.batch_size,
                                 collate_fn=collate_fn,
                                 drop_last=True,
                                 shuffle=True)
    return train_dataloader, test_dataloader


if __name__ == '__main__':
    # nd = NerDataset(datas)
    # print(nd.__len__())
    get_data()
